package nachos.vm;

import nachos.machine.*;
import nachos.threads.*;
import nachos.userprog.*;

import java.util.HashMap;
import java.util.LinkedList;

/**
 * A <tt>UserProcess</tt> that supports demand-paging.
 */
public class VMProcess extends UserProcess
{
    /**
     * Allocate a new process.
     */
    public VMProcess()
    {
        super();
        kernel = (VMKernel) ThreadedKernel.kernel;
    }

    /**
     * Most systems simply invalidate all the TLB entries on a context switch,
     * causing entries to be reloaded as the pages are referenced.
     *
     * Save the state of this process in preparation for a context switch.
     * Called by <tt>UThread.saveState()</tt>.
     */
    public void saveState()
    {
        // invalidate all TLB entries
        kernel.propagateAndFlushTLB(true);
    }

    /**
     * Restore the state of this process after a context switch. Called by
     * <tt>UThread.restoreState()</tt>.
     */
    public void restoreState()
    {
        // 重写这个方法，避免父类setPageTable的操作
    }

    /**
     * Initializes page tables for this process so that the executable can be
     * demand-paged.
     *
     * @return    <tt>true</tt> if successful.
     */
    protected boolean loadSections()
    {
        int topVPN = 0;
        for (int i = 0; i < coff.getNumSections(); i++) {
            CoffSection section = coff.getSection(i);

            CoffConstructor coffConstructor;

            topVPN += section.getLength();
            for (int j = section.getFirstVPN(); j < topVPN; j++) {
                coffConstructor = new CoffConstructor(section, j);
                sectionsMap.put(j, coffConstructor);
            }

        }

        for (; topVPN < numPages - 1; topVPN++) {
            sectionsMap.put(topVPN, new StackConstructor(topVPN));
        }

        return true;
    }

    /**
     * Release any resources allocated by <tt>loadSections()</tt>.
     */
    protected void unloadSections()
    {
        kernel.freePages(PID, numPages);
    }

    protected void loadArguments(int entryOffset, int stringOffset, byte[][] argv) {
        sectionsMap.put(numPages - 1, new ArgConstructor(entryOffset, stringOffset, argv));
    }

    /**
     * Handle a user exception. Called by
     * <tt>UserKernel.exceptionHandler()</tt>. The
     * <i>cause</i> argument identifies which exception occurred; see the
     * <tt>Processor.exceptionZZZ</tt> constants.
     *
     * @param    cause    the user exception that occurred.
     */
    public void handleException(int cause)
    {
        Processor processor = Machine.processor();

        switch (cause)
        {
            case Processor.exceptionTLBMiss:
                handleTLBMiss(processor.readRegister(Processor.regBadVAddr));
                break;
            default:
                super.handleException(cause);
                break;
        }
    }

    private void handleTLBMiss(int vaddr) {
        if (!validAddress(vaddr)) {
            return;
        }
        TranslationEntry retrievedTranslationEntry = retrievePage(Processor.pageFromAddress(vaddr));

        boolean unwritten = true;
        for (int i = 0; i < Machine.processor().getTLBSize() && unwritten; i++) {
            TranslationEntry entry = Machine.processor().readTLBEntry(i);

            if (entry.ppn == retrievedTranslationEntry.ppn) {
                Machine.processor().writeTLBEntry(i, retrievedTranslationEntry);
                unwritten = false;
            } else if (!entry.valid) {
                Machine.processor().writeTLBEntry(i, retrievedTranslationEntry);
                unwritten = false;
            }
        }

        // 采用随机替换
        if (unwritten) {
            int randomIndex = Lib.random(Machine.processor().getTLBSize());
            TranslationEntry oldEntry = Machine.processor().readTLBEntry(randomIndex);

            if (oldEntry.dirty || oldEntry.used) {
                kernel.propagateEntry(oldEntry.ppn, oldEntry.used, oldEntry.dirty);
            }

            Machine.processor().writeTLBEntry(randomIndex, retrievedTranslationEntry);
        }

        kernel.unpin(retrievedTranslationEntry.ppn);
    }

    public TranslationEntry retrievePage(int vpn) {
        TranslationEntry entry;
        if (sectionsMap.containsKey(vpn)) {
            entry = sectionsMap.get(vpn).execute();
        } else if ((entry = kernel.pinIfExists(vpn, PID)) == null) {
            entry = kernel.pageFault(vpn, PID);
        }

        Lib.assertTrue(entry != null);
        return entry;
    }


    @Override
    public int readVirtualMemory(int vaddr, byte[] data, int offset, int length) {
        int bytesRead = 0;
        LinkedList<VMMemoryAccess> memoryAccesses = createMemoryAccesses(vaddr, data, offset, length, AccessType.READ, true);

        //Execute them if they were successfully created
        if (memoryAccesses != null) {
            int temp;
            for (VMMemoryAccess vma : memoryAccesses) {
                temp = vma.executeAccess();
                if (temp == 0)
                    break;
                else
                    bytesRead += temp;
            }
        }

        return bytesRead;
    }

    @Override
    public int writeVirtualMemory(int vaddr, byte[] data, int offset, int length) {
        return writeVirtualMemory(vaddr, data, offset, length, true);
    }

    public int writeVirtualMemory(int vaddr, byte[] data, int offset, int length, boolean unpin) {
        int bytesWritten = 0;
        LinkedList<VMMemoryAccess> memoryAccesses = createMemoryAccesses(vaddr, data, offset, length, AccessType.WRITE, unpin);

        //Execute them if they were successfully created
        if (memoryAccesses != null) {
            int temp;
            for (VMMemoryAccess vma : memoryAccesses) {
                temp = vma.executeAccess();
                if (temp == 0)
                    break;
                else
                    bytesWritten += temp;
            }
        }

        return bytesWritten;
    }

    public int writeVirtualMemory(int vaddr, byte[] data, boolean unpin) {
        return VMProcess.this.writeVirtualMemory(vaddr, data, 0, data.length, unpin);
    }

    public LinkedList<VMMemoryAccess> createMemoryAccesses(int vaddr,
                                                           byte[] data,
                                                           int offset,
                                                           int length,
                                                           AccessType accessType,
                                                           boolean unpin) {
        Lib.assertTrue(offset >= 0 && length >= 0 && offset + length <= data.length);
        LinkedList<VMMemoryAccess> returnList  = null;

        if (validAddress(vaddr)) {
            returnList = new LinkedList<VMMemoryAccess>();

            while (length > 0) {
                int vpn = Processor.pageFromAddress(vaddr);

                int potentialPageAccess = Processor.pageSize - Processor.offsetFromAddress(vaddr);
                int accessSize = length < potentialPageAccess ? length : potentialPageAccess;

                returnList.add(new VMMemoryAccess(accessType, data, vpn, offset, Processor.offsetFromAddress(vaddr), accessSize, unpin));
                length -= accessSize;
                vaddr += accessSize;
                offset += accessSize;
            }
        }

        return returnList;
    }

    private static final int pageSize = Processor.pageSize;
    private static final char dbgProcess = 'a';
    private static final char dbgVM = 'v';

    private VMKernel kernel;

    public HashMap<Integer, Constructor> sectionsMap = new HashMap<>();


    // --------------------------------------------
    // Inner classes
    // --------------------------------------------


    protected class VMMemoryAccess extends UserProcess.MemoryAccess {
        VMMemoryAccess(AccessType accessType, byte[] data, int vpn, int dataStart, int pageStart, int length, boolean unpin) {
            super(accessType, data, vpn, dataStart, pageStart, length);
            this.unpin = unpin;
        }

        @Override
        public int executeAccess() {
            //Overwrite the translationEntry with a new one
            //page should be pinned
            translationEntry = retrievePage(vpn);

            int bytesAccessed = super.executeAccess();

            if (unpin) {
                kernel.unpin(translationEntry.ppn);
            }

            return bytesAccessed;
        }

        /** A bit to indicate whether the access should unpin the page when it is finished. */
        public boolean unpin;
    }

    public interface Constructor {
        TranslationEntry execute();
    }

    public class CoffConstructor implements Constructor {
        CoffConstructor(CoffSection section, int vpn) {
            this.vpn = vpn;
            this.coffSection = section;
        }

        @Override
        public TranslationEntry execute() {
            int sectionNum = vpn - coffSection.getFirstVPN();
            Lib.assertTrue(sectionsMap.remove(vpn) != null);

            // 获取一个可用页
            TranslationEntry entry = kernel.requestFreePage(vpn, PID);
            coffSection.loadPage(sectionNum, entry.ppn);

            entry.readOnly = coffSection.isReadOnly();
            return entry;
        }

        int vpn;
        CoffSection coffSection;
    }

    public class StackConstructor implements Constructor {
        StackConstructor(int vpn) {
            this.vpn = vpn;
        }

        @Override
        public TranslationEntry execute() {
            Lib.assertTrue(sectionsMap.remove(vpn) != null);

            TranslationEntry entry = kernel.requestFreePage(vpn, PID);
            entry.readOnly = false;
            return entry;
        }

        int vpn;
    }

    public class ArgConstructor implements Constructor {
        public ArgConstructor(int entryOffset, int stringOffset, byte[][] argv) {
            this.entryOffset = entryOffset;
            this.stringOffset = stringOffset;
            this.argv = argv;
        }

        @Override
        public TranslationEntry execute() {
            Lib.assertTrue(sectionsMap.remove(numPages - 1) != null);

            TranslationEntry entry = kernel.requestFreePage(numPages - 1, PID);
            for (int i = 0; i < argv.length; i++) {
                byte[] stringOffsetBytes = Lib.bytesFromInt(stringOffset);
                Lib.assertTrue(writeVirtualMemory(entryOffset, stringOffsetBytes) == 4);
                entryOffset += 4;
                Lib.assertTrue(writeVirtualMemory(stringOffset, argv[i]) == argv[i].length);
                stringOffset += argv[i].length;
                Lib.assertTrue(writeVirtualMemory(stringOffset, new byte[] {0}) == 1);
                stringOffset++;
            }

            entry.readOnly = true;
            return entry;
        }

        public int entryOffset;
        public int stringOffset;
        public byte[][] argv;
    }
}
